package utils

import (
	"context"
	"fmt"
	"time"
)

type LimitedResource[T any] struct {
	maxResources    int
	createFunc      func() (T, error)
	destroyFunc     func(T) error
	resourceChannel chan *T // 资源通道使用指针类型
	timeout         time.Duration
}

func NewLimitedResource[T any](
	maxResources int,
	createFunc func() (T, error),
	destroyFunc func(T) error,
	timeout time.Duration,
) *LimitedResource[T] {
	lr := &LimitedResource[T]{
		maxResources:    maxResources,
		createFunc:      createFunc,
		destroyFunc:     destroyFunc,
		resourceChannel: make(chan *T, maxResources),
		timeout:         timeout,
	}

	if err := lr.Initialize(); err != nil {
		panic(err) // Ensure resources are initialized or panic
	}

	return lr
}

func (lr *LimitedResource[T]) Borrow(ctx context.Context) (*T, error) {
	select {
	case resource := <-lr.resourceChannel:
		// Create a new context for resource timeout
		resourceCtx, cancel := context.WithTimeout(ctx, lr.timeout)
		go func(res *T) {
			select {
			case <-resourceCtx.Done():
				// Context timeout occurred
				fmt.Println("auto destory")
				lr.Return(res) // Return resource to channel or destroy if channel full
			}
			cancel() // Clean up context
		}(resource)
		return resource, nil
	case <-ctx.Done():
		return nil, ctx.Err() // Context canceled or deadline exceeded
	}
}

func (lr *LimitedResource[T]) Return(resource *T) error {
	select {
	case lr.resourceChannel <- resource:
		return nil
	default:
		// Channel is full, destroy resource
		return lr.destroyFunc(*resource)
	}
}

func (lr *LimitedResource[T]) Initialize() error {
	for i := 0; i < lr.maxResources; i++ {
		resource, err := lr.createFunc()
		if err != nil {
			return err
		}
		resourcePtr := &resource
		lr.resourceChannel <- resourcePtr
	}
	return nil
}

func (lr *LimitedResource[T]) Shutdown() error {
	close(lr.resourceChannel)

	var err error
	for resource := range lr.resourceChannel {
		if destroyErr := lr.destroyFunc(*resource); destroyErr != nil {
			err = destroyErr
		}
	}
	return err
}
